import shutil
import images
import os
from timer import timer
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as pp
import math
from kmeans import kmeans
from hist import hist

__dtype__ = 'float32'
__gendata__ = True
__csv__ = True
__draw__ = False
__print__ = False

__drawcsv__ = True

class gmm:
    data = None
    data_files = None
    __dbg__ = False
    __history__ = False
    hist_mean = None
    hist_sigma = None
    hist_weight = None
    means = None
    sigmas = None
    weights = None
    G = None
    runs = None
    r = None
    __timer__ = None
    count, bins, patches = None, None, None
    aic, aicc = None, None
    neg_log_lik = None
    neg_hist, pos_hist = None, None
    neg_bins, pos_bins = None, None
    limit = None
    
    def __init__(self, data_files=None, debug=False, history=False, timer=False, texts=False, gmm_txt=None, hist_txt=None, limit=None):
        ''' create a new model with specified data '''

        if texts:
            if debug:
                print 'Creating GMM from precompiled data'

            gmm_data = np.load(gmm_txt)
            self.means = gmm_data['means']
            self.G = len(self.means)
            self.sigmas = gmm_data['sigmas']
            self.weights = gmm_data['weights']
            self.aic = gmm_data['aic']
            # self.aicc = gmm_data['aicc']
            
            print 'txt', hist_txt
            hist_data = np.load(hist_txt)
            self.bins = hist_data['bins']
            self.count = hist_data['count']
        else:        
            if data_files is not None:
                self.data_files = data_files
                if debug:
                    print 'Creating GMM with', len(data_files), 'data files'
            self.__dbg__ = debug
            self.__history__ = history
            self.__timer__ = timer
 	    self.limit = limit

    @staticmethod
    def gendata(n, mean, sigma):
        ''' generate n data points from a specified pdf '''
        return np.random.randn(n).astype(__dtype__) * sigma + mean

    @staticmethod
    def random_data(collection):
        ''' generate data based on collection, which is a list of tuples (n, mean, sigma) '''
        data = np.array([], dtype=__dtype__)
        for entry in collection:
            n, mean, sigma = entry
            data = np.concatenate((data, gmm.gendata(n, mean, sigma)))
        return data

    @staticmethod
    def random_poisson(L):
        ''' generate data based on collection, which is a list of tuples (n, mean, sigma) '''
        data = np.random.poisson(l, 1000000)
        return data

    @staticmethod
    def random_lognorm(sigma=1, mu=0):
        ''' generate samples from log normal distribution with provided parameters '''
        data = np.random.lognormal(mu, sigma, 100000)
        return data

    @staticmethod
    def not_zero(x):
        ''' make sure x is not zero before returning '''
        if x == 0:
            return 1e-100
        else:
            return x

    def P(self):
        ''' probability that a value x was generated from a Gaussian distr with sigma and mu '''
        sqrtpi = math.sqrt(2.*math.pi)
        '''
        neg = np.array([(1./(self.sigmas[i] * sqrtpi) * np.exp(-0.5*(self.neg_bins-self.means[i])**2./self.sigmas[i]**2.)) for i in range(self.G)], dtype='float32')
        pos = np.array([(1./(self.sigmas[i] * sqrtpi) * np.exp(-0.5*(self.pos_bins-self.means[i])**2./self.sigmas[i]**2.)) for i in range(self.G)], dtype='float32')
        neg[np.where(neg == 0)] = 1e-5
        pos[np.where(pos == 0)] = 1e-5
        return neg, pos
        '''
        output = np.array([(1./(self.sigmas[i] * sqrtpi) * np.exp(-0.5*(self.data-self.means[i])**2./self.sigmas[i]**2.)) for i in range(self.G)], dtype='float32')
        output[np.where(output == 0)] = 1e-10
        return output
        
    @staticmethod    
    def p(x, mu, sigma):
        ''' probability that a value x was generated from a Gaussian distr with sigma and mu '''
        return (1./(sigma * math.sqrt(2.*math.pi))) * math.exp(-0.5*(x-mu)**2./sigma**2.)

    def GMM_P(self, PX):
        ''' joint probability that x is generated by the gmm, with p(x) precalculated '''
        '''
        PROB = [np.dot(self.weights, PX[i]).astype(__dtype__) for i in range(2)]

        for i in range(2):
            PROB[i][np.where(PROB == 0)] = 1e-4
        '''
        PROB = np.dot(self.weights, PX).astype(__dtype__)

        PROB[np.where(PROB == 0)] = 1e-4
        return PROB

    def gmm_p(self, x):
        ''' joint probability that x is generated by the gmm with means, sigmas and weights '''
        prob = 0
        for i in range(len(self.means)):
            prob += gmm.p(x, self.means[i], self.sigmas[i]) * self.weights[i]
            
        return prob

    def PIX(self, PX, GMMX):
        ''' p(i | x) probability of x being generated from all G pdf '''
        self.weights = np.array(self.weights, dtype=__dtype__)

        '''
        weighted = [np.multiply(np.transpose(np.atleast_2d(self.weights)), PX[i], dtype=__dtype__) for i in range(2)]
        pix = [np.array([np.divide(weighted[i], GMMX[i])], dtype=__dtype__) for i in range(2)]
        pix = [pix[i][0, :, :] for i in range(2)]

        return pix, [np.sum(weighted[i], axis=0, dtype=__dtype__) for i in range(2)]
        '''

        weighted = np.multiply(np.transpose(np.atleast_2d(self.weights)), PX, dtype=__dtype__)
        pix = np.array([np.divide(weighted, GMMX)], dtype=__dtype__)
        pix = pix[0, :, :]

        return pix, np.sum(weighted, axis=0, dtype=__dtype__)
    
    def pix(self, x, i):
        ''' p(i | x) probability of x being generated from this pdf '''
        return (gmm.p(x, self.means[i], self.sigmas[i]) * self.weights[i]) / gmm.not_zero(self.gmm_p(x))

    def fit_model(self, G, total_runs=100, tol=1e1):
        ''' fit a gmm to the data supplied by running an EM algorithm '''
        another = True

        self.G = G

        self.runs = 0
        self.neg_log_lik = 0
        logs = list()

        # preprocess input
        # if self.__dbg__:
        #     print 'precalculating histogram'
        # self.neg_hist, self.neg_bins, self.pos_hist, self.pos_bins = hist(self.data)

        # rerun until convergence
        while another:
            another = False

            self.weights = self.G * [1./G]
            if self.__dbg__:
                print 'starting kmeans'
            # self.means = np.random.choice(self.data, G)
            km = kmeans(self.data_files[-1], self.G, debug=self.__dbg__)
            self.means = km.means
            self.sigmas = self.G * [2.5]
            print 'km:', km.means
            km = None
            
            if self.__history__:
                self.hist_mean = np.zeros((total_runs, self.G))
                self.hist_sigma = np.zeros((total_runs, self.G))
                self.hist_weight = np.zeros((total_runs, self.G))
            
            if self.__dbg__:            
                print 'begin conditions:'
                print self.weights
                print self.means
                print self.sigmas
            
            for self.r in range(total_runs):
                if self.__dbg__:
                    print "Performing iteration", self.r
                    
                br = True

                #for i in range(G):
                oldmeans = self.means
                oldsigmas = self.sigmas
                oldweights = self.weights
                oldloglik = self.neg_log_lik

                try:
                    newmeans, newsigmas, newweights, negloglik = self.update_values()
                    print newmeans, negloglik
                except ZeroDivisionError:
                    if self.__dbg__:
                        print "ZERODIVISIONERROR"
                        print "ANOTHER!!"
                    another = True
                    br = True
                    break
                
                if np.any(np.isnan(newmeans)) or np.any(np.isnan(newsigmas)) or np.any(np.isnan(newweights)):
                    br = True
                    another = True
                    # raise Exception("Gayyy")
                    if self.__dbg__:
                        print "ANOTHER!!"
                    break

                '''
                # check for empty components
                if np.size(np.where(newweights == 0)[0]) > 0:
                    another = True
                    br = True
                    if self.__dbg__:
                        print "EMPTY COMPONENT"
                        print "ANOTHER!!"
                    break
                '''

                self.means = newmeans
                self.sigmas = newsigmas
                self.weights = newweights
                self.neg_log_lik = negloglik
                logs.append(negloglik)
                
                if self.__history__:
                    self.hist_mean[self.r] = newmeans
                    self.hist_sigma[self.r] = newsigmas
                    self.hist_weight[self.r] = newweights

                #if np.all(np.abs(newsigmas - oldsigmas) < tol) and np.all(np.abs(newmeans - oldmeans) < tol) and np.all(np.abs(newweights - oldweights) < tol) and br:
                if self.__dbg__:
                    print '\t\t\tloglik', negloglik
                if abs(negloglik - oldloglik) < tol and br:
                    break
                # else:
                #br = False

        k = len(self.weights[np.where(self.weights != 0)])
        self.aic = 2. * k + 2. * self.neg_log_lik
        self.aicc = self.aic + (2*k*(k+1))/(len(self.data) - k - 1)

    @staticmethod
    def read_data(fn):
        dat_file = np.load(fn)
        return dat_file['data']
            
    def update_values(self):
        ''' update mean, sigma and weight by using EM '''
        k = 0
        top_weights = np.zeros((self.G,), dtype=__dtype__)
        top_means = np.zeros((self.G,), dtype=__dtype__)
        top_sigmas = np.zeros((self.G,), dtype=__dtype__)
        neg_log_lik = 0.
        N = 0
        for fn in self.data_files:
            # load next data set
            self.data = gmm.read_data(fn)
	    if self.limit is not None:
		self.data = self.data[np.where(np.abs(self.data) < self.limit)]
            #print np.mean(self.data), np.sqrt(np.var(self.data))
            #print np.min(self.data), np.max(self.data)
            N += len(self.data)
            
            # precalculate all p(x) and gmm(x)
            tm = timer()
            if self.__dbg__:
                print 'calculating px'
            PX = self.P()
            if self.__dbg__:
                print 'precalc px', tm.elapsed()
            tm.reset()

            if self.__dbg__:
                print 'calculating gmmx'
            GMMX = self.GMM_P(PX)
            if self.__dbg__:
                print 'precalc gmmx', tm.elapsed()

            tm.reset()
            if self.__dbg__:
                print 'calculating pix'
            PIX, weighted = self.PIX(PX, GMMX)
            if self.__dbg__:
                print 'precalc pix', tm.elapsed()

            top_weights += np.sum(PIX, axis=1)
            top_means   += np.array([np.dot(PIX[i], self.data) for i in range(self.G)])
            top_sigmas  += np.array([np.dot(PIX[i], np.square(np.abs(self.data - self.means[i]))) for i in range(self.G)])
            neg_log_lik += -np.sum(np.log(weighted))

        '''
        # multiply in histogram
        PIX = [PIX[0] * self.neg_hist, PIX[1] * self.pos_hist]

        top_weights = np.sum(PIX[0], axis=1) + np.sum(PIX[1], axis=1)
        top_means   = np.array([np.dot(PIX[0][i], self.neg_bins) for i in range(self.G)]) + np.array([np.dot(PIX[1][i], self.pos_bins) for i in range(self.G)])

        top_sigmas  = np.array([np.dot(PIX[0][i], np.square(np.abs(self.neg_bins - self.means[i]))) for i in range(self.G)]) + np.array([np.dot(PIX[1][i], np.square(np.abs(self.pos_bins - self.means[i]))) for i in range(self.G)])
        neg_log_lik = -np.sum(np.log(weighted[0])) -np.sum(np.log(weighted[1]))
        '''
        
        # top_mean = top_sigma = top_weight = 0.
        #if self.__timer__:
        #    #tm = timer(0, len(self.data))
        #    #print 'starting update_values'
            
        # for j in range(len(self.data)):
        #   # px = self.pix(self.data[j], i)
        #    top_mean += px * self.data[j]
        #    top_sigma += px * (abs(self.data[j] - self.means[i]) ** 2.)
        #    top_weight += px
        #    #if self.__timer__:
        #    #    tm.progress(j)

        # mean, sigma, weight
        return np.divide(top_means, top_weights), np.sqrt(np.divide(top_sigmas, top_weights)), top_weights / N, neg_log_lik

    def calc_hist(self, bins=None):
        ''' helper function to calculate hist without drawing or saving '''
        N = 0
        if self.__dbg__:
            print 'Creating histogram'
        if bins is None:
            for fn in self.data_files:
                self.data = gmm.read_data(fn)
		if self.limit is not None: self.data = self.data[np.where(np.abs(self.data) < self.limit)]
                count, bins = np.histogram(self.data, bins=1000, normed=False) #, weights = np.ones_like(self.data)/len(self.data))
                if self.bins is None:
                    self.bins = bins
                if self.count is None:
                    self.count = count
                else:
                    self.count += count
                N += len(self.data) 
        else:
            for fn in self.data_files:
                self.data = gmm.read_data(fn)
		if self.limit is not None: self.data = self.data[np.where(np.abs(self.data) < self.limit)]
                count, bins = np.histogram(self.data, bins=bins, normed=False)#, weights = np.ones_like(self.data)/float(len(self.data)))
                if self.bins is None:
                    self.bins = bins
                if self.count is None:
                    self.count = count
                else:
                    self.count += count
                N += len(self.data)
        
        self.count = self.count.astype('float')
        self.count /= N
        
    def save_results(self, library, alg, hist=False):
        ''' save the results of the gmm and hist to a file for later processing '''
        ''' formats:
            means
            sigmas
            weights

            and:
            count
            bins
            '''

        print 'saving %s/gmm_%s_%d.npz' % (library, alg, self.G)
        try:
            os.mkdir(library)
        except OSError:
            pass
        np.savez_compressed('%s/gmm_%s_%d.npz' % (library, alg, self.G), means=self.means, sigmas=self.sigmas, weights=self.weights, aic=self.aic)
	shutil.copyfile('%s/gmm_%s_%d.npz' % (library, alg, self.G), '/copy/gmm_%s_%d.npz' % (alg, self.G))

        if hist:
	    print 'saving %s/hist_%s_%d.npz' % (library, alg, self.G)
            if self.count is None or self.bins is None:
                print 'calc hist'
                self.calc_hist(bins=range(-150,151))
            np.savez_compressed('%s/hist_%s_%d.npz' % (library, alg, self.G), count=self.count, bins=self.bins)
	    shutil.copyfile('%s/hist_%s_%d.npz' % (library, alg, self.G), '/copy/hist_%s_%d.npz' % (alg, self.G))

    def print_results(self):
        print 'means   =', self.means
        print 'sigmas  =', self.sigmas
        print 'weights =', self.weights
        print 'aic     =', self.aic
        print
        
    def draw_hist(self, bins=None, fig=True, limit_hist=True, xlim=None):
        ''' draw only the calculated histogram '''

        if len(self.bins) == 0 or self.bins is None or self.count is None:
            print 'calculating hist'
            self.calc_hist(bins)
        # self.count, self.bins, self.patches = pp.hist(self.data, bins=bins, normed=False, weights = np.ones_like(self.data)/len(self.data))

        
        #fig, ax = pp.subplots()
        
        self.bins = self.bins[:-1]
        
        self.bins = np.array(self.bins).astype(__dtype__)
        self.count = np.array(self.count).astype(__dtype__)

        self.bins = np.reshape(self.bins, (np.size(self.bins), 1))
        self.count = np.reshape(self.count, (np.size(self.count), 1))
        
        if fig: pp.figure()
        
        pp.bar(self.bins - (self.bins[1] - self.bins[0])/2., self.count, width=1.0)
        if xlim is not None:
            pp.xlim(xlim)
        
        if limit_hist:
            tmp = list()
            tmpc = list()
            for i in range(len(self.bins)):
                if self.count[i] != 0:
                    break
            zero = True
    
            for j in range(len(self.bins)-1, i, -1):
                if zero == True:
                    if self.count[j] != 0:
                        zero = False
    
                if zero == False:
                    tmp.append(self.bins[j])
                    tmpc.append(self.count[j])
    
            self.bins = np.array(tmp)[::-1]
            print 'bins', self.bins
            self.count = np.array(tmpc)[::-1]

    def draw_gmm(self, fig=False):
        ''' draw only the calculated gmm '''
        
        if fig: pp.figure()

        # PDF
        x = 0.1 * np.array(range(int((np.min(self.bins)) - 3) * 10, int(np.max(self.bins)) * 10))
        y = np.zeros_like(x, dtype='float')

        for i in range(self.G):
            for j in range(len(x)):
                y[j] += self.weights[i] * gmm.p(x[j], self.means[i], self.sigmas[i])

        pp.hold()
        pp.plot(x, y, c='r', linewidth=3)

    def draw_gmm_hist(self, fig=True):
        ''' draw the discrete form of the gmm '''

        if fig: pp.figure()

        # PDF
        x = self.bins
        bin_width = x[1] - x[0]
        y = np.zeros_like(x, dtype='float')
        
        # for i in range(self.G):
        for j in range(len(x)):
            y[j] = self.pdf_integral(x[j] - bin_width/2., x[j] + bin_width/2.)

        # x += bin_width/2.
        
        pp.hold()
        # pp.stem(x, y, markerfmt='xr', linefmt='r-')
        pp.plot(x, y, color='r', linewidth=3)
        
    def cdf(self, x):
        return np.sum([self.weights[i] * 0.5 * (1 + math.erf((float(x) - self.means[i]) / (math.sqrt(2.0) * self.sigmas[i]))) for i in range(self.G)])
    
    def draw_cdf(self, fig=True):
        # CDF
        if fig: pp.figure()
        z = np.zeros_like(x, dtype='float')

        #for i in range(self.G):
        for j in range(len(x)):
            z[j] = self.cdf(x[j])

        pp.plot(x, z)

    def pdf_integral(self, x, y):
        ''' return the integral of the pdf given by this model from x to y '''
        srt = np.sort([x, y])
        return self.cdf(srt[1]) - self.cdf(srt[0])
    
    '''
    def draw_results(self, bins=range(-100,102)):
        '''''' plot the resulting gmm ''''''
        pp.cla()
        pp.close('all')
        if self.__dbg__:
            print 'Drawing results'
        
        pp.figure(1)
        pp.subplot(1,3,1)
        pp.plot(range(self.r), self.hist_mean[:self.r, :])
        pp.legend(range(self.G))
        pp.title('Means')
        pp.subplot(1,3,2)
        pp.plot(range(self.r), self.hist_sigma[:self.r, :])
        pp.legend(range(self.G)) 
        pp.title('Sigmas')
        pp.subplot(1,3,3)
        pp.plot(range(self.r), self.hist_weight[:self.r, :])
        pp.ylim([0, 1])
        pp.legend(range(self.G))
        pp.title('Weights')
        
        self.draw_hist()
        self.draw_gmm()
     '''   

def save_csv(csv_data):
    ''' save aic data to csv for later drawing '''
    fi = file('aic.csv', 'w')
    for dat in csv_data:
        fi.write('%.2f,%.3f,%.3f,%.3f,%.3f\n' % (dat[0], dat[1], dat[2], dat[3], dat[4]))
        

if __name__ == '__main__':
    if __drawcsv__:
        # pp.figure()
        
        sigmas = list()
        y = list()
        for line in file('aic.csv'):
            print line
            sigma, one, two, three, four = line.split(',')
            if float(sigma) < 6:
                if float(sigma) != 0.75:
                    sigmas.append(float(sigma))
                    y.append(np.array([float(one), float(two), float(three), float(four)]) / float(one))
            
        sigmas = ['\$\s = %.2f\$' % sig for sig in sigmas]
        
        images.save_latex(range(1,5), np.array(y).T, sigmas, 'gmm_test', xlabel='GMM components (G)', ylabel='Normalised AIC', width=0.7, draw=False, ylim=[0.4, 1.1], yticks=[0.4, 0.6, 0.8, 1.0])
        
    else:
        tm = timer(0,1)
        G = range(1, 5)
        # L = range(1, 2)
        SIGMA = np.array(range(1,5)) / 4.
        SIGMA = [SIGMA, np.array(range(2, 10, 2))]
        SIGMA = np.array(SIGMA)
        SIGMA = SIGMA[np.where(SIGMA > 0)]
        
        if __gendata__:
            print 'generating data'
            # for l in L:
            for sig in SIGMA:
                # data = gmm.random_poisson(L)
                data = gmm.random_lognorm(sig)
                # np.savez('poisson_%d.npz' % l, data=data)
                np.savez('lognorm_%.2f.npz' % sig, data=data)
                
        # data = gmm.random_data([(100000, 0, 2.5), (75000, 25, 0.5), (75000, 5, 0.5)])
    #     data = np.round(data).astype('int')
    #     np.savez('data.npz', data=data)
    
        if __csv__: csv = list()
        for sig in SIGMA:
            if __csv__: 
                aics = list()
                aics.append(0)
            for g in G:
                # model = gmm(['data.npz', 'data.npz'], debug=False, history=True)
                model = gmm(['lognorm_%.2f.npz' % sig], debug=False, history=True)
                model.fit_model(g)
                total_runs = 1000
        
                model.save_results('lognorm', '%.2f' % sig, hist=True)
                print 'time', tm.elapsed()
                
                # '%s/gmm_%s_%d.npz' % (library, alg, self.G)
                model2 = gmm(texts = True, gmm_txt='lognorm/gmm_%.2f_%d.npz' % (sig, g), hist_txt='lognorm/hist_%.2f_%d.npz' % (sig, g))
                print 
                print 'read model (sigma = %.2f, G = %d):' % (sig, g)
                if __print__: model2.print_results()
                if __draw__:
                    model2.draw_hist()
                    pp.hold()
                    model2.draw_gmm(fig=False)
                if __csv__: aics.append(float(model2.aic))
            
            if __csv__:
                aics[0] = sig
                csv.append(aics)
                # model2.draw_gmm_hist(fig=False)
                # model2.draw_cdf()
        # model2.draw_hist()
        
        if __csv__: save_csv(csv)
        if __draw__: pp.show()
        # model.draw_results()
    
    print 'done'
    # pp.show()
